#include "hypercube/hypercube_cryto_agent.hpp"
#include "hypercube/hypercube_hmm_engine.hpp"
#include "matrix_util.hpp"
#include <gtest/gtest.h>

static params<helib::BGV> my_params;

#define SCALE 1000.0
#define INV_SCALE 1.0 / (SCALE)

#define HMM_FP_TEST(DIM, IB)                                                   \
  TEST(TEST_HMM_FP, MNK_##DIM) {                                               \
    ASSERT_LE(IB + IB + std::log2(DIM),                                        \
              static_cast<double>(my_params.t_bits));                          \
                                                                               \
    unsigned long m = DIM;                                                     \
    unsigned long n = DIM;                                                     \
    unsigned long k = DIM;                                                     \
                                                                               \
    double *A = nullptr;                                                       \
    double *B = nullptr;                                                       \
    double *C = nullptr;                                                       \
    std::size_t size_A = m * k;                                                \
    std::size_t size_B = k * n;                                                \
    std::size_t size_C = m * n;                                                \
    ALLOC_DOUBLE_ARRAY_WITH_INIT(A, size_A);                                   \
    ALLOC_DOUBLE_ARRAY_WITH_INIT(B, size_B);                                   \
    ALLOC_DOUBLE_ARRAY_WITH_INIT(C, size_C);                                   \
                                                                               \
    FIX_RAND_SEED(false);                                                      \
    RAND_ARRAY(A, size_A, IB);                                                 \
    RAND_ARRAY(B, size_B, IB);                                                 \
                                                                               \
    NTL::mat_ZZ ptxt_A;                                                        \
    NTL::mat_ZZ ptxt_B;                                                        \
    ptxt_A.SetDims(m, k);                                                      \
    ptxt_B.SetDims(k, n);                                                      \
    COPY_ARRAY_TO_NTL_MATRIX(ptxt_A, A);                                       \
    COPY_ARRAY_TO_NTL_MATRIX(ptxt_B, B);                                       \
                                                                               \
    double scale = 1.0;                                                        \
    SCALE_DOUBLE_ARRAY(A, size_A, INV_SCALE);                                  \
    scale *= INV_SCALE;                                                        \
    SCALE_DOUBLE_ARRAY(B, size_B, INV_SCALE);                                  \
    scale *= INV_SCALE;                                                        \
    cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, m, n, k, 1.0, A, k, \
                B, n, 0.0, C, n);                                              \
    /*std::stringstream ss;                                                    \
    ss << std::fixed << std::setprecision(6);                                  \
    PRINTF_ARRAY_AS_MATRIX(ss, C, m, n);                                       \
    std::cout << ss.str() << std::endl;*/                                      \
                                                                               \
    const auto eh_cube = hypercube_hmme::get_hypercube(m, n, k);               \
    hypercube_hmmcc hhmmcc(my_params, eh_cube, true);                          \
                                                                               \
    hmm_status<shmm_engine> status(METHOD::BSGS, omp_get_max_threads());       \
    hypercube_hmme hhmme = hypercube_hmme(status);                             \
    auto eh_mnk = hypercube_hmme::get_expected_mnk(m, n, k);                   \
    hhmme.register_engine(hhmmcc.context()->getEA(), eh_mnk);                  \
                                                                               \
    auto ctxt_A = hhmmcc.encrypt(ptxt_A);                                      \
    auto ctxt_B = hhmmcc.encrypt(ptxt_B);                                      \
                                                                               \
    auto ctxt_C = hhmmcc.get_empty_ctxt();                                     \
    auto type = shmm_engine::parse_matrix_type(eh_mnk);                        \
    hhmme.cgemm(type, ctxt_C, ctxt_A, ctxt_B);                                 \
    NTL::mat_ZZ ptxt_C;                                                        \
    hhmmcc.decrypt(ptxt_C, ctxt_C);                                            \
    /*std::cout << ptxt_C << std::endl;*/                                      \
                                                                               \
    for (auto i = 0; i < m; i++) {                                             \
      for (auto j = 0; j < n; j++) {                                           \
        auto expected = C[i * n + j];                                          \
        auto actual = NTL::conv<double>(ptxt_C[i][j]) * scale;                 \
        /* EXPECT_EQ(expected, actual); */                                     \
        ASSERT_LE(std::abs(expected - actual), scale)                          \
            << " diff at index (" << i << ", " << j << ")";                    \
      }                                                                        \
    }                                                                          \
    delete[] A;                                                                \
    delete[] B;                                                                \
    delete[] C;                                                                \
  }

HMM_FP_TEST(32, 10)
HMM_FP_TEST(64, 10)
HMM_FP_TEST(128, 10)

int main(int argc, char **argv) {
  ::testing::InitGoogleTest(&argc, argv);

  command_line_parser parser(my_params);
  int ret = parser.parse(argc, argv);
  EXPECT_EQ(ret, 0);

  // Run tests
  ret = RUN_ALL_TESTS();

  return ret;
}